
package com.waleipt.api.es;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

import org.apache.commons.lang3.RandomUtils;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
import org.elasticsearch.threadpool.ThreadPool;
import org.elasticsearch.threadpool.ThreadPoolStats;
import org.junit.Test;

import junit.framework.Assert;

public class ThreadPoolTest {


	protected final ThreadPool.Info info(final ThreadPool threadPool, final String name) {
		for (final ThreadPool.Info info : threadPool.info()) {
			if (info.getName().equals(name)) {
				return info;
			}
		}	
		assert "same".equals(name);
		return null;
	}

	protected final ThreadPoolStats.Stats stats(final ThreadPool threadPool, final String name) {
		for (final ThreadPoolStats.Stats stats : threadPool.stats()) {
			if (name.equals(stats.getName())) {
				return stats;
			}
		}
		throw new IllegalArgumentException(name);
	}

	protected final void terminateThreadPoolIfNeeded(final ThreadPool threadPool) throws InterruptedException {
		if (threadPool != null) {
			terminate(threadPool);
		}
	}

	public static boolean terminate(ThreadPool service) throws InterruptedException {
		return ThreadPool.terminate(service, 10, TimeUnit.SECONDS);
	}

	static String randomThreadPool(final ThreadPool.ThreadPoolType type) {
		List<String> list = ThreadPool.THREAD_POOL_TYPES.entrySet().stream().filter(t -> t.getValue().equals(type))
				.map(Map.Entry::getKey).collect(Collectors.toList());
		return list.get(RandomUtils.nextInt(0, list.size() - 1));
	}

	@Test
	public void testSimple()
	{
		final Settings nodeSettings =
	            Settings.builder()
	                .put("node.name", "testSimple")
	                .put("thread_pool." + ThreadPool.Names.SAME + ".size", 10)
	                .put("thread_pool." + ThreadPool.Names.SAME + ".queue_size", 10)
	                .build();
		ThreadPool threadPool = new ThreadPool( nodeSettings );
		final CountDownLatch latch = new CountDownLatch( 100 );
		final AtomicInteger counter = new AtomicInteger(0) ;
		for ( int i =0; i< 100; i++ )
		{
			threadPool.executor( ThreadPool.Names.SAME ).execute( () -> {
				try
				{
					counter.addAndGet(1) ;
				}
				finally
				{
					latch.countDown(); 
				}
			});
		}
		
		try {
			latch.await();
		} catch (InterruptedException e) {
			e.printStackTrace();
		} 
		
		System.err.println( counter.get() );
		
	}
	
	
	@Test
	public void testThreadPool() {

		final String threadPoolName = randomThreadPool(ThreadPool.ThreadPoolType.FIXED);
        // some of the fixed thread pool are bound by the number of
        // cores so we can not exceed that
        final int size = 4;
        final int queueSize = RandomUtils.nextInt(1, 16);
        final long rejections = RandomUtils.nextInt(1, 16);

        ThreadPool threadPool = null;
        final Settings nodeSettings =
            Settings.builder()
                .put("node.name", "testRejectedExecutionCounter")
                .put("thread_pool." + threadPoolName + ".size", size)
                .put("thread_pool." + threadPoolName + ".queue_size", queueSize)
                .build();
        try {
            threadPool = new ThreadPool(nodeSettings);

            // these tasks will consume the thread pool causing further
            // submissions to queue
            final CountDownLatch latch = new CountDownLatch(size);
            final CountDownLatch block = new CountDownLatch(1);
            for (int i = 0; i < size; i++) {
                threadPool.executor(threadPoolName).execute(() -> {
                    try {
                        latch.countDown();
                        block.await();
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                });
            }

            // wait for the submitted tasks to be consumed by the thread
            // pool
            try {
				latch.await();
			} catch (InterruptedException e1) {
				e1.printStackTrace();
			}

            // these tasks will fill the thread pool queue
            for (int i = 0; i < queueSize; i++) {
                threadPool.executor(threadPoolName).execute(() -> {});
            }

            // these tasks will be rejected
            long counter = 0;
            for (int i = 0; i < rejections; i++) {
                try {
                    threadPool.executor(threadPoolName).execute(() -> {});
                } catch (EsRejectedExecutionException e) {
                    counter++;
                }
            }

            block.countDown();

            Assert.assertTrue( (counter==rejections) );
//            assertThat(stats(threadPool, threadPoolName).getRejected(), equalTo(rejections));
        } finally {
//            terminateThreadPoolIfNeeded(threadPool);
        }
	}

}
